import numpy as np
import torch
import torch.utils.data as Data
from torch.autograd import Variable

from .evaluation import Evaluation
from .evaluation import MIN_COMPENSATION

class Ncoverage(Evaluation):
    def __init__(self, outputs_origin,outputs_adv, device, **kwargs):
        '''
        @description:
        @param {
            model:
            device:
            kwargs:
        }
        @return: None
        '''
        super(Ncoverage, self).__init__(outputs_origin,outputs_adv, device)

        self._parsing_parameters(**kwargs)

    def _parsing_parameters(self, **kwargs):
        '''
        @description:
        @param {
        }
        @return:
        '''

    def evaluate(self,adv_xs=None, cln_xs=None, cln_ys=None,adv_ys=None,target_preds=None, target_flag=False):
        '''
        @description:
        @param {
            adv_xs: 攻击样本
            cln_xs：原始样本
            cln_ys: 原始类别，非目标攻击下原始样本的类型
            adv_ys: 攻击样本的预测类别
            target_preds： 目标攻击下希望原始样本攻击的目标类别
            target_flag：是否是目标攻击
        }
        @return: Neuron Coverage
        '''
        total = len(adv_xs)
        print("total",total)
        outputs = self.outputs_adv
        assert len(adv_xs) == len(cln_ys), 'examples and labels do not match.'
        coverage = 0.0
        for i in range(outputs.shape[0]):
            pred_i = outputs[i].data.cpu().numpy()
            coverage_i = pred_i >= pred_i[cln_ys[i]]
            coverage_i = coverage_i.sum() / len(pred_i)
            coverage += coverage_i
        number = adv_xs.shape[0]
        if not number==0:
            neuron_coverage = coverage / number
        else:
            neuron_coverage = coverage / (number+MIN_COMPENSATION)
        return neuron_coverage